import time
import requests
from utils import get_client
import os


def check_batch_status(batch_id, output_dir="batch_results"):
    """检查批量任务状态并下载结果"""
    client = get_client("SILICONFLOW")
    os.makedirs(output_dir, exist_ok=True)

    while True:
        batch = client.batches.retrieve(batch_id)
        status = batch.status
        print(f"任务状态: {status}")

        if status == 'completed':
            file_id = batch.output_file_id
            output_path = os.path.join(output_dir, f"batch_{batch_id}_results.jsonl")

            # 下载文件到本地
            response = requests.get(file_id, stream=True)

            if response.status_code == 200:
                with open(output_path, 'wb') as f:
                    for chunk in response.iter_content(chunk_size=8192):
                        if chunk:
                            f.write(chunk)
                print(f"文件下载成功: {output_path}")
            else:
                print(f"文件下载失败，状态码: {response.status_code}")

            return output_path

        elif status == 'failed':
            raise Exception("批量任务失败")

        time.sleep(30)  # 每30秒检查一次